package nz.ac.massey.cs.plot;

import static java.lang.Math.max;
import static java.lang.Math.min;
import static nz.ac.massey.cs.stats.utils.Numeric.maximum;
import static nz.ac.massey.cs.stats.utils.Numeric.minimum;

import java.awt.Font;
import java.awt.Graphics;
import java.awt.Graphics2D;
import java.awt.Polygon;
import java.awt.Shape;
import java.awt.geom.Ellipse2D;
import java.awt.geom.Line2D;
import java.awt.geom.Rectangle2D;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import nz.ac.massey.cs.reg.LinearRegression;
import nz.ac.massey.cs.reg.Regressor;

/**
 * @author Hussain Al-Mutawa
 * @version 3.0
 */
public class ScatterPlot extends AbstractPlot {

	private double xScale, yScale;
	private boolean showRegressionLine, showEquation, showOriginLines;
	private LinkedList<Double> xs;
	private LinkedList<Double> ys;

	public ScatterPlot(List<Double> xs, List<Double> ys) {
		super();

		this.xs = new LinkedList<Double>(xs);
		this.ys = new LinkedList<Double>(ys);
		while (this.ys.size() > this.xs.size()) {
			this.ys.removeLast();
		}
		while (this.xs.size() > this.ys.size()) {
			this.xs.removeLast();
		}
		setXmin(max(0, minimum(xs) - 0.1));
		setYmin(max(0, minimum(ys) - 0.1));
		setXmax(maximum(xs) + 0.1);
		setYmax(maximum(ys) + 0.1);
	}

	public ScatterPlot() {
		this(new LinkedList<Double>(),
				new LinkedList<Double>());
	}

	public void add(double x, double y) {
		xs.add(x);
		ys.add(y);
		setXmin(min(getXmin(), x));
		setXmax(max(getXmax(), x));
		setYmin(min(getYmin(), y));
		setYmax(max(getYmax(), y));
	}

	@Override
	public void redraw(Graphics graphics) {
		Graphics2D g = (Graphics2D) graphics;

		setxScale((getxAxisEnd() - getxAxisStart()) / (getXmax() - getXmin()));
		setyScale((getyAxisEnd() - getyAxisStart()) / (getYmax() - getYmin()));

		drawBackground(g);
		g.setStroke(getNormalStroke());
		drawXAxis(g);
		drawXLabel(g);
		drawXIntervals(g);
		drawYAxis(g);
		drawYLabel(g);
		drawYIntervals(g);
		drawXYPlots(g, xs, ys);
		if (isShowRegressionLine()) {
			Regressor reg = drawRegressionLine(g);
			if (isShowEquation()) {
				drawEquation(g, reg);
			}
		}
		if (isShowOriginLines()) {
			drawOriginLines(g);
		}
		drawTitle(g);
	}

	public void drawOriginLines(Graphics2D g) {
		double y0 = getyAxisEnd() + getYmin() * getyScale();
		double x0 = getxAxisStart() - getXmin() * getxScale();

		g.setColor(getOriginAxisColor());
		g.setStroke(getOriginStroke());

		if (y0 > getyAxisStart() && y0 < getyAxisEnd()) {
			g.draw(new Line2D.Double(getxAxisStart() - 4, y0, getxAxisEnd(), y0));
		}
		if (x0 > getxAxisStart() && x0 < getxAxisEnd()) {
			g.draw(new Line2D.Double(x0, getyAxisStart(), x0, getyAxisEnd() + 4));
		}
	}

	public void drawEquation(Graphics2D g, Regressor reg) {
		g.setStroke(getOriginStroke());
		g.setColor(getEquationColor());
		g.setFont(new Font(g.getFont().getFontName(), Font.BOLD, 12));
		String str = String.format("y = %1.2fx + %1.2f", reg.beta().get(1), reg.beta().get(0));
		int len = g.getFontMetrics().stringWidth(str);
		g.drawString(str, (float) getxAxisEnd() - len, (float) getxAxisStart() - 5);
	}

	public Regressor drawRegressionLine(Graphics2D g) {
		Map<Double, Double> XY = new HashMap<Double, Double>();
		for (int i = 0; i < xs.size(); i++) {
			if (XY.containsKey(xs.get(i))) {
				XY.put(xs.get(i), (XY.get(xs.get(i)) + ys.get(i)) / 2);
			} else {
				XY.put(xs.get(i), ys.get(i));
			}
		}
		Regressor reg = new LinearRegression(XY);

		double x1 = getxAxisStart();
		double x2 = getxAxisStart() + ((-getXmin() + getXmax()) * getxScale());
		double y1 = getyAxisEnd() - ((-getYmin() + reg.predict(getXmin())) * getyScale());
		double y2 = getyAxisEnd() - ((-getYmin() + reg.predict(getXmax())) * getyScale());

		g.setStroke(getTitleStroke());
		g.setColor(getRegressionLineColor());
		g.draw(new Line2D.Double(x1, y1, x2, y2));

		return reg;
	}

	public double getScaledX(double x) {
		return -5 + getxAxisStart() + ((-getXmin() + x) * getxScale());
	}

	public double getScaledY(double y) {
		return -5 + getyAxisEnd() - ((-getYmin() + y) * getyScale());
	}

	public void drawXYPlots(Graphics2D g, List<Double> xs, List<Double> ys) {

		double x, y;

		g.setColor(getPlotColor());
		g.setFont(getPlotFont());

		for (int i = 0; i < xs.size(); i++) {
			x = getScaledX(xs.get(i));
			y = getScaledY(ys.get(i));
			plot(g,x,y);
		}
	}

	public void plot(Graphics2D g,double x,double y){
		int ix, iy;
		iy = (int) y;
		ix = (int) x;

		switch (getPlotChar()) {
		case '.':
			g.fill(new Ellipse2D.Double(x, y, 3, 3));
			break;
		case 'o':
			g.draw(new Ellipse2D.Double(x, y, 5, 5));
			break;
		case '@':
			g.fill(new Ellipse2D.Double(x, y, 5, 5));
			break;
		case 'D':
			g.draw(new Rectangle2D.Double(x, y, 5, 5));
			break;
		case 'x':
			g.drawString("x", ix - 2, iy - 2);
			break;
		default:
			Shape polygon = new Polygon(
					new int[]{ix, ix - 3, ix + 3},
					new int[]{iy - 3, iy + 3, iy + 3},
					3);
			g.draw(polygon);
			break;
		}
	}

	public boolean isShowRegressionLine() {
		return showRegressionLine;
	}

	public void setShowRegressionLine(boolean showRegressionLine) {
		this.showRegressionLine = showRegressionLine;
	}

	public boolean isShowEquation() {
		return showEquation;
	}

	public void setShowEquation(boolean showEquation) {
		this.showEquation = showEquation;
	}

	public double getxScale() {
		return xScale;
	}

	public double getyScale() {
		return yScale;
	}

	public void setxScale(double xScale) {
		this.xScale = xScale;
	}

	public void setyScale(double yScale) {
		this.yScale = yScale;
	}

	public boolean isShowOriginLines() {
		return showOriginLines;
	}

	@Override
	public String toString() {
		return "Scatterplot\n" + xs + "\n" + ys;
	}


}